package com.gy.ai.agent.client;

import com.gy.ai.agent.listener.ReqClientListener;
import com.gy.ai.agent.util.SysOkHttpClientUtil;
import com.gy.ai.agent.util.SysUtil;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import okio.ByteString;
import org.jetbrains.annotations.NotNull;

import java.util.Map;

/**
 * @author LGY
 * Create by 2025/6/18 15:53
 */
@Slf4j
public class WebSocketClient extends WebSocketListener implements BaseReqClient {

    private WebSocket webSocket;
    private ReqClientListener listener;
    private boolean isStop = false;
    private String payload;

    @Override
    public void start(String url, Map<String, String> headers, String payload, ReqClientListener listener) {
        this.listener = listener;
        this.payload = payload;
        OkHttpClient client = SysOkHttpClientUtil.buildDefaultClient();
        Request request = new Request.Builder()
                .url(url)
                .build();
        this.isStop = false;
        this.webSocket = client.newWebSocket(request, this);
        log.info("消息发送载体: {}", payload);
    }

    @Override
    public void stop() {
        try {
            tryToStop();
        } finally {
            tryToCloseWebSocket();
        }
    }


    //webSocket events
    @Override
    public void onOpen(WebSocket webSocket, Response response) {
        webSocket.send(payload);
        this.listener.onStart(this);
    }

    @Override
    public void onMessage(WebSocket webSocket, String text) {
        log.info("消息接收载体：{}", text);
        this.listener.onMessage(this, text);
    }

    @Override
    public void onMessage(WebSocket webSocket, ByteString bytes) {
        this.onMessage(webSocket, bytes.utf8());
    }

    @Override
    public void onClosing(WebSocket webSocket, int code, String reason) {
        try {
            tryToStop();
        } finally {
            tryToCloseWebSocket();
        }
    }

    @Override
    public void onFailure(WebSocket webSocket, Throwable t, Response response) {
        try {
            try {
                Throwable failureThrowable = SysUtil.getFailureThrowable(t, response);
                this.listener.onFailure(this, failureThrowable);
            } finally {
                tryToStop();
            }
        } finally {
            tryToCloseWebSocket();
        }
    }

    @Override
    public void onClosed(@NotNull WebSocket webSocket, int code, @NotNull String reason) {
        try {
            tryToStop();
        } finally {
            tryToCloseWebSocket();
        }
    }


    private void tryToCloseWebSocket() {
        if (this.webSocket != null) {
            this.webSocket.close(1000, "");
            this.webSocket = null;
        }
    }

    private boolean tryToStop() {
        if (!this.isStop) {
            this.isStop = true;
            this.listener.onStop(this);
            return true;
        }
        return false;
    }
}
